import argparse
from collections import defaultdict

def parse_arguments():
    parser = argparse.ArgumentParser(description="Sort a TSV file by specific columns with memory efficiency.")
    parser.add_argument("input", help="Input TSV file")
    parser.add_argument("output", help="Output sorted TSV file")
    parser.add_argument("--delimiter", default="\t", help="Column delimiter (default: tab)")
    parser.add_argument("--col3", type=int, default=3, help="Column for primary sorting (default: 3)")
    parser.add_argument("--col1", type=int, default=1, help="Column for secondary sorting (default: 1)")
    parser.add_argument("--col9", type=int, default=9, help="Column for tertiary sorting (default: 9)")
    return parser.parse_args()

def record_sort_keys(input_file, delimiter, col3, col1, col9):
    sort_keys = []
    key_count = defaultdict(int)  # 用于统计每个排序键的出现次数
    line_data = []

    with open(input_file, "r") as f:
        for line_no, line in enumerate(f):
            columns = line.strip().split(delimiter)
            if len(columns)<col9:
                print(f"Line {line_no + 1} has less than {col9} columns. Skipping.")
                continue
            key3 = columns[col3 - 1]
            key1 = columns[col1 - 1]
            try:
                key9 = float(columns[col9 - 1])
            except ValueError:
                print(f"Line {line_no + 1} has an invalid value in column {col9}. Skipping.")
                key9 = float('inf')  # 如果第9列的值不是有效数字，处理为无穷大
                continue
            key = (key3, key1, key9)
            key_count[key] += 1
            line_data.append((key, line_no))  # 记录每行的排序键和行号

    return line_data, key_count

def sort_by_keys(line_data):
    # 按照排序键（key3, key1, key9）对行数据进行排序
    line_data.sort(key=lambda x: (x[0][0], x[0][1], x[0][2]))
    return line_data

def output_sorted_file(input_file, output_file, sorted_keys, key_count, delimiter):
    with open(input_file, "r") as infile, open(output_file, "w") as outfile:
        # 读取文件的所有行
        lines = list(infile)
        
        for key, line_no in sorted_keys:
            if key_count[key] >= 5:  # 只输出出现次数大于或等于5的行
                outfile.write(lines[line_no])

def main():
    args = parse_arguments()
    
    # 第一遍扫描，记录排序键和行号，并统计排序键出现的次数
    line_data, key_count = record_sort_keys(args.input, args.delimiter, args.col3, args.col1, args.col9)
    
    # 第二遍扫描，排序并过滤掉出现次数少于5的行
    sorted_keys = sort_by_keys(line_data)
    
    # 输出排序后的文件
    output_sorted_file(args.input, args.output, sorted_keys, key_count, args.delimiter)

if __name__ == "__main__":
    main()
